"""
This script is used to make data augmentation, 
the method is very easy which is just adding white noise to the raw data.
"""
import numpy as np
import pandas as pd
import os
from constant import TRAIN_PATH, DATASET_PATH
from tqdm import tqdm


def main():
    add_zero_mean_gaussian()


def add_zero_mean_gaussian():
    dataset = pd.read_csv(DATASET_PATH)
    dataset = dataset[dataset.category != 2]  # delete the old augmentation items
    dataset_events = dataset[dataset.category == 1].copy()  # 1 means events data
    dataset_columns = dataset.columns
    augm_path = os.path.join(TRAIN_PATH, "augmentation")

    process_bar = tqdm(total=22 * dataset_events.shape[0])
    for factor in range(0, 22):
        for row in dataset_events.itertuples():
            data_path = row.data_path
            file_name = os.path.basename(data_path)
            process_bar.set_description(f"{factor}, {file_name}")
            process_bar.update(1)
            batch_data = np.load(data_path)
            for i in range(batch_data.shape[0]):
                data = batch_data[i]
                data_std = np.std(data, axis=1)
                for j in range(data_std.shape[0]):
                    # The way I deal with the data augmentation is a little different from the author
                    noise = np.random.normal(loc=0, scale=0.1 * factor * data_std[j], size=data.shape[1])
                    data[j] += noise
                batch_data[i] = data

            save_name = file_name.split('.')[0]
            save_path = os.path.join(augm_path, f'{save_name}_std_{factor}.npy')
            np.save(save_path, batch_data)
            df = pd.DataFrame(data=[[row.data_info_path, save_path, row.amount, 2]], columns=dataset_columns)
            dataset = pd.concat([dataset, df], ignore_index=True)
    dataset.to_csv(DATASET_PATH, index=False)
    process_bar.close()


if __name__ == "__main__":
    main()
